#ifndef DISPLAY_H
#define DISPLAY_H

#include <iostream>
#include <vector>
#include <Eigen/Core>

#include <ros/ros.h>
#include <visualization_msgs/MarkerArray.h>
#include <nav_msgs/Path.h>
#include "quad_msgs/Trajectory.h"
#include "quad_msgs/ReplanCheck.h"

namespace Planner
{
    using Eigen::MatrixXd;
    using Eigen::Vector3d;
    using std::cout;
    using std::endl;
    using std::ios_base;
    using std::string;
    using std::vector;

    // 打印矩阵
    void printMatrix(const MatrixXd &Q)
    {

        ios_base::fmtflags origin = cout.setf(ios_base::fixed, ios_base::floatfield);
        std::streamsize prec = cout.precision(2);
        for (int i = 0; i < Q.rows(); i++)
        {
            for (int j = 0; j < Q.cols(); j++)
                cout << Q(i, j) << "\t";
            cout << "\n";
        }
        cout.setf(origin, ios_base::floatfield);
        cout.precision(prec);
    }

    void printMatrix(const vector<MatrixXd> &Q)
    {

        ios_base::fmtflags origin = cout.setf(ios_base::fixed, ios_base::floatfield);
        std::streamsize prec = cout.precision(2);
        for (int k = 0; k < Q.size(); k++)
            printMatrix(Q[k]);
        cout.setf(origin, ios_base::floatfield);
        cout.precision(prec);
    }

    template <typename T>
    void printVector(const vector<T> &v)
    {
        ios_base::fmtflags origin = cout.setf(ios_base::fixed, ios_base::floatfield);
        std::streamsize prec = cout.precision(2);
        for (size_t i = 0; i < v.size(); i++)
            cout << v[i] << " ";
        cout << "\n";
        cout.setf(origin, ios_base::floatfield);
        cout.precision(prec);
    }

    void printVector(const vector<Vector3d> &v)
    {
        ios_base::fmtflags origin = cout.setf(ios_base::fixed, ios_base::floatfield);
        std::streamsize prec = cout.precision(2);
        for (int i = 0; i < v.size(); i++)
            cout << v[i][0] << " " << v[i][1] << " " << v[i][2] << endl;
        cout.setf(origin, ios_base::floatfield);
        cout.precision(prec);
    }

    void displayWaypoints(const vector<Vector3d> &waypoints, ros::Publisher &publisher, string frame = "world")
    {
        visualization_msgs::MarkerArray pts;
        visualization_msgs::Marker pt;

        pt.header.frame_id = frame;
        pt.type = visualization_msgs::Marker::SPHERE;
        pt.scale.x = 0.2;
        pt.scale.y = 0.2;
        pt.scale.z = 0.2;
        pt.color.a = 1.f;
        pt.color.r = 1.f;
        pt.color.g = 0.f;
        pt.color.b = 0.f;

        for (int i = 0; i < waypoints.size(); i++)
        {
            pt.id = i;
            pt.pose.position.x = waypoints[i][0];
            pt.pose.position.y = waypoints[i][1];
            pt.pose.position.z = waypoints[i][2];
            pts.markers.push_back(pt);
        }
        publisher.publish(pts);
    }

    void displayTrajectory(const vector<Vector3d> &waypoints, ros::Publisher &publisher, string frame = "world")
    {
        nav_msgs::Path traj_msgs;
        traj_msgs.header.frame_id = frame;
        geometry_msgs::PoseStamped pt_msgs;
        pt_msgs.header.frame_id = frame;

        for (auto p : waypoints)
        {
            pt_msgs.pose.position.x = p[0];
            pt_msgs.pose.position.y = p[1];
            pt_msgs.pose.position.z = p[2];
            traj_msgs.poses.push_back(pt_msgs);
        }
        publisher.publish(traj_msgs);
    }

    void displayTrajectory(const vector<MatrixXd> &coeffs, const vector<double> &time_seq,
                           ros::Publisher &publisher, double gap = 0.01, string frame = "world")
    {
        nav_msgs::Path traj_msgs;
        traj_msgs.header.frame_id = frame;
        geometry_msgs::PoseStamped pt_msgs;
        pt_msgs.header.frame_id = frame;

        int segments = time_seq.size();
        int order = coeffs[0].cols();

        for (int s = 0; s < segments; s++)
        {
            double duration = time_seq[s];
            for (double t = 0; t < duration; t += gap)
            {
                double x = 0, y = 0, z = 0;
                for (int r = 0; r < order; r++)
                {
                    x += coeffs[0](s, r) * pow(t, r);
                    y += coeffs[1](s, r) * pow(t, r);
                    z += coeffs[2](s, r) * pow(t, r);
                }
                pt_msgs.pose.position.x = x;
                pt_msgs.pose.position.y = y;
                pt_msgs.pose.position.z = z;
                traj_msgs.poses.push_back(pt_msgs);
            }
        }
        publisher.publish(traj_msgs);
    }

    void displayTrajectory(const quad_msgs::Trajectory &msg, ros::Publisher &publisher,
                           double gap = 0.1, string frame = "world")
    {
        nav_msgs::Path traj_msgs;
        traj_msgs.header.frame_id = frame;
        geometry_msgs::PoseStamped pt_msgs;
        pt_msgs.header.frame_id = frame;

        for (int s = 0; s < msg.segments; s++)
        {
            double duration = msg.polynomials[s].duration;
            for (double t = 0; t < duration; t += gap)
            {
                double x = 0, y = 0, z = 0;
                for (int r = 0; r < msg.order; r++)
                {
                    x += msg.polynomials[s].coeffs_x[r] * pow(t, r);
                    y += msg.polynomials[s].coeffs_y[r] * pow(t, r);
                    z += msg.polynomials[s].coeffs_z[r] * pow(t, r);
                }
                pt_msgs.pose.position.x = x;
                pt_msgs.pose.position.y = y;
                pt_msgs.pose.position.z = z;
                traj_msgs.poses.push_back(pt_msgs);
            }
        }
        publisher.publish(traj_msgs);
    }

    void displayTrajectory(const quad_msgs::ReplanCheck::Request &request, ros::Publisher &publisher,
                           double gap = 0.1, string frame = "world")
    {
        nav_msgs::Path traj_msgs;
        traj_msgs.header.frame_id = frame;
        geometry_msgs::PoseStamped pt_msgs;
        pt_msgs.header.frame_id = frame;

        int order = request.polynomials[0].coeffs_x.size();
        for (int s = 0; s < request.polynomials.size(); s++)
        {
            double duration = request.polynomials[s].duration;
            for (double t = 0; t < duration; t += gap)
            {
                double x = 0, y = 0, z = 0;
                for (int r = 0; r < order; r++)
                {
                    x += request.polynomials[s].coeffs_x[r] * pow(t, r);
                    y += request.polynomials[s].coeffs_y[r] * pow(t, r);
                    z += request.polynomials[s].coeffs_z[r] * pow(t, r);
                }
                pt_msgs.pose.position.x = x;
                pt_msgs.pose.position.y = y;
                pt_msgs.pose.position.z = z;
                traj_msgs.poses.push_back(pt_msgs);
            }
        }
        publisher.publish(traj_msgs);
    }

    void displayTrajectory(const quad_msgs::ReplanCheck::Response &response, ros::Publisher &publisher,
                           double gap = 0.1, string frame = "world")
    {
        if (response.replan_flag == quad_msgs::ReplanCheck::Response::SAFE ||
            response.replan_flag == quad_msgs::ReplanCheck::Response::REPLAN_FAIL)
            return;
        
        nav_msgs::Path traj_msgs;
        traj_msgs.header.frame_id = frame;
        geometry_msgs::PoseStamped pt_msgs;
        pt_msgs.header.frame_id = frame;

        int order = response.polynomials[0].coeffs_x.size();
        for (int s = 0; s < response.polynomials.size(); s++)
        {
            double duration = response.polynomials[s].duration;
            for (double t = 0; t < duration; t += gap)
            {
                double x = 0, y = 0, z = 0;
                for (int r = 0; r < order; r++)
                {
                    x += response.polynomials[s].coeffs_x[r] * pow(t, r);
                    y += response.polynomials[s].coeffs_y[r] * pow(t, r);
                    z += response.polynomials[s].coeffs_z[r] * pow(t, r);
                }
                pt_msgs.pose.position.x = x;
                pt_msgs.pose.position.y = y;
                pt_msgs.pose.position.z = z;
                traj_msgs.poses.push_back(pt_msgs);
            }
        }
        publisher.publish(traj_msgs);
    }

    void printTrajectoryMsg(const quad_msgs::Trajectory &msg)
    {
        cout << "Trajectory information:\n";
        cout << "order: " << msg.order << ", segments: " << msg.segments << endl;
        cout << "waypoint 0: " << msg.waypoints[0].position.x << " "
             << msg.waypoints[0].position.y << " " << msg.waypoints[0].position.z << endl;
        for (int i = 0; i < msg.segments; i++)
        {
            cout << "duration: " << msg.polynomials[i].duration << endl;
            cout << "waypoint " << i + 1 << ": " << msg.waypoints[i + 1].position.x << " "
                 << msg.waypoints[i + 1].position.y << " " << msg.waypoints[i + 1].position.z << endl;
        }
    }

} // namespace Planner

#endif